/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 *  The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.spark.sql.execution.datasources.greenplum

import java.sql.{Connection, PreparedStatement}

import org.apache.spark.Partition
import org.apache.spark.rdd.RDD
import org.apache.spark.sql._
import org.apache.spark.sql.execution.datasources.jdbc.{JDBCRDD, JdbcUtils}
import org.apache.spark.sql.jdbc.{JdbcDialect, JdbcDialects, JdbcType}
import org.apache.spark.sql.sources.{BaseRelation, Filter, InsertableRelation, PrunedFilteredScan}
import org.apache.spark.sql.types._


private[sql] case class GreenplumRelation(
    parts: Array[Partition], options: GreenplumOptions)(@transient val sparkSession: SparkSession)
  extends BaseRelation
  with PrunedFilteredScan
  with InsertableRelation {

  import GreenplumUtils._

  override def sqlContext: SQLContext = sparkSession.sqlContext

  override val needConversion: Boolean = false
  override val schema: StructType = {
    val tableSchema = JDBCRDD.resolveTable(options)
    options.customSchema match {
      case Some(customSchema) => JdbcUtils.getCustomSchema(
        tableSchema, customSchema, sparkSession.sessionState.conf.resolver)
      case None => tableSchema
    }
  }

  // Check if JDBCRDD.compileFilter can accept input filters
  override def unhandledFilters(filters: Array[Filter]): Array[Filter] = {
    filters.filter(JDBCRDD.compileFilter(_, JdbcDialects.get(options.url)).isEmpty)
  }

  override def buildScan(requiredColumns: Array[String], filters: Array[Filter]): RDD[Row] = {
    // Rely on a type erasure hack to pass RDD[InternalRow] back as RDD[Row]
    JDBCRDD.scanTable(
      sparkSession.sparkContext,
      schema,
      requiredColumns,
      filters,
      parts,
      options).asInstanceOf[RDD[Row]]
  }

  override def insert(data: DataFrame, overwrite: Boolean): Unit = {
    val conn = JdbcUtils.createConnectionFactory(options)()
    val maxConns = options.maxConnections
    try {
      if (overwrite) {
        if (options.isTruncate &&
          JdbcUtils.isCascadingTruncateTable(options.url).contains(false)) {
          JdbcUtils.truncateTable(conn, options)
          nonTransactionalCopy(
            if (options.transactionOn) data.coalesce(1) else data.coalesce(maxConns),
            schema, options)
        } else {
          transactionalCopy(data.coalesce(maxConns), schema, options)
        }
      } else {
        nonTransactionalCopy(
          if (options.transactionOn) {
            data.coalesce(1)
          } else {
            data.coalesce(maxConns)
          },
          schema, options)
      }
    } finally {
      closeConnSilent(conn)
    }
  }

  private def getJdbcType(dt: DataType, dialect: JdbcDialect): JdbcType = {
    dialect.getJDBCType(dt).getOrElse(  // orElse(makeConverter(dt))
      throw new IllegalArgumentException(s"Can't get JDBC type for ${dt.catalogString}"))
  }

  private type JDBCValueSetter_add = (PreparedStatement, Row, Int,Int) => Unit

  private def makeSetter(conn: Connection, dataType: DataType): JDBCValueSetter_add = dataType match {
    case IntegerType =>
      (stmt: PreparedStatement, row: Row, pos: Int,currentpos:Int) =>
        stmt.setInt(pos + 1, row.getInt(pos-currentpos))

    case LongType =>
      (stmt: PreparedStatement, row: Row, pos: Int,currentpos:Int) =>
        stmt.setLong(pos + 1, row.getLong(pos-currentpos))

    case DoubleType =>
      (stmt: PreparedStatement, row: Row, pos: Int,currentpos:Int) =>
        stmt.setDouble(pos + 1, row.getDouble(pos-currentpos))

    case FloatType =>
      (stmt: PreparedStatement, row: Row, pos: Int,currentpos:Int) =>
        stmt.setFloat(pos + 1, row.getFloat(pos-currentpos))

    case ShortType =>
      (stmt: PreparedStatement, row: Row, pos: Int,currentpos:Int) =>
        stmt.setInt(pos + 1, row.getShort(pos-currentpos))

    case ByteType =>
      (stmt: PreparedStatement, row: Row, pos: Int,currentpos:Int) =>
        stmt.setInt(pos + 1, row.getByte(pos-currentpos))

    case BooleanType =>
      (stmt: PreparedStatement, row: Row, pos: Int,currentpos:Int) =>
        stmt.setBoolean(pos + 1, row.getBoolean(pos-currentpos))

    case StringType =>
      (stmt: PreparedStatement, row: Row, pos: Int,currentpos:Int) =>
        stmt.setString(pos + 1, row.getString(pos-currentpos))

    case BinaryType =>
      (stmt: PreparedStatement, row: Row, pos: Int,currentpos:Int) =>
        stmt.setBytes(pos + 1, row.getAs[Array[Byte]](pos-currentpos))

    case TimestampType =>
      (stmt: PreparedStatement, row: Row, pos: Int,currentpos:Int) =>
        stmt.setTimestamp(pos + 1, row.getAs[java.sql.Timestamp](pos-currentpos))

    case DateType =>
      (stmt: PreparedStatement, row: Row, pos: Int,currentpos:Int) =>
        stmt.setDate(pos + 1, row.getAs[java.sql.Date](pos-currentpos))

    case t: DecimalType =>
      (stmt: PreparedStatement, row: Row, pos: Int,currentpos:Int) =>
        stmt.setBigDecimal(pos + 1, row.getDecimal(pos-currentpos))
    case _ =>
      (stmt: PreparedStatement, row: Row, pos: Int,currentpos:Int) =>
        throw new IllegalArgumentException(
          s"Can't translate non-null value for field $pos")
  }

  /**
    * 批量插入 或更新 数据 ,该方法 借鉴Spark.write.save() 源码
    * @param dataFrame 数据框
    * @param sc sc
    * @param table
    * @param id
    */
  def insertOrUpdateToPgsql(dataFrame:DataFrame,table:String,id:String): Unit ={

    val tableSchema = dataFrame.schema
    val columns =tableSchema.fields.map(x => x.name).mkString(",")
    val placeholders = tableSchema.fields.map(_ => "?").mkString(",")
    val sql = s"INSERT INTO $table ($columns) VALUES ($placeholders) on conflict($id) do update set "
    val update = tableSchema.fields.map(x =>
      x.name.toString + "=?"
    ).mkString(",")

    val realsql =sql.concat(update)
    val conn = JdbcUtils.createConnectionFactory(options)() // connectionPool()
    conn.setAutoCommit(false)
    val dialect = JdbcDialects.get(conn.getMetaData.getURL)
    val broad_ps = sparkSession.sparkContext.broadcast(conn.prepareStatement(realsql))

    val numFields = tableSchema.fields.length *2
    //调用spark中自带的函数 或者 捞出来,获取属性字段与字段类型
    val nullTypes = tableSchema.fields.map(f => getJdbcType(f.dataType, dialect).jdbcNullType)
    val setters = tableSchema.fields.map(f => makeSetter(conn,f.dataType))

    var rowCount = 0
    val batchSize = 2000
    val updateindex = numFields / 2
    try {
      dataFrame.foreachPartition(iterator =>{
        //遍历批量提交
        val ps = broad_ps.value
        try{
          while (iterator.hasNext) {
            val row = iterator.next()
            var i = 0
            while (i < numFields) {
              i < updateindex match {
                case true =>{
                  if (row.isNullAt(i)) {
                    ps.setNull(i + 1, nullTypes(i))
                  } else {
                    setters(i).apply(ps, row, i,0)
                  }
                }
                case false =>{
                  if (row.isNullAt(i-updateindex)) {
                    ps.setNull(i + 1, nullTypes(i-updateindex))
                  } else {
                    setters(i-updateindex).apply(ps, row, i,updateindex)
                  }
                }
              }
              i = i + 1
            }
            ps.addBatch()
            rowCount += 1
            if (rowCount % batchSize == 0) {
              ps.executeBatch()
              rowCount = 0
            }
          }
          if (rowCount > 0) {
            ps.executeBatch()
          }
        }finally {
          ps.close()
        }
      })
      conn.commit()
    }catch {
      case e: Exception =>
//        logError("Error in execution of insert. " + e.getMessage)
        conn.rollback()
      // insertError(connectionPool("OuCloud_ODS"),"insertOrUpdateToPgsql",e.getMessage)
    }finally {
      conn.close()
    }
  }


  override def toString: String = {
    val partitioningInfo = if (parts.nonEmpty) s" [numPartitions=${parts.length}]" else ""
    // credentials should not be included in the plan output, table information is sufficient.
    s"GreenplumRelation(${options.table})" + partitioningInfo
  }
}
